import numpy as np
import pandas as pd


def calculate_features(data):
    """计算技术指标特征（优化版）"""
    # 价格指标
    if 'ts_code' in data.columns:
        df = data.drop(columns=['ts_code'])
    data['MA5'] = data['close'].rolling(5).mean()
    data['MA20'] = data['close'].rolling(20).mean()
    data['MA_diff'] = data['MA5'] - data['MA20']

    # 波动率指标
    data['ATR'] = data['high'].rolling(14).max() - data['low'].rolling(14).min()

    # 成交量指标
    data['VOL_MA10'] = data['volume'].rolling(10).mean()
    data['VOL_MA20'] = data['volume'].rolling(20).mean()

    # 动量指标（优化RSI计算）
    delta = data['close'].diff()
    gain = delta.where(delta > 0, 0)
    loss = -delta.where(delta < 0, 0)
    avg_gain = gain.rolling(14).mean()
    avg_loss = loss.rolling(14).mean()
    data['RSI'] = 100 - (100 / (1 + (avg_gain / avg_loss.replace(0, np.nan))))

    # 滞后特征（过去3天收盘价）
    for lag in [1, 2, 3]:
        data[f'close_lag{lag}'] = data['close'].shift(lag)

    return data.dropna(),['MA5', 'MA20', 'MA_diff', 'ATR', 'VOL_MA10', 'VOL_MA20', 'RSI',
            'close_lag1', 'close_lag2', 'close_lag3']


def add_technical_indicators(df):
    """
    专业级均线偏离特征处理器（含鲁棒性错误处理）
    改进点：动态窗口验证、稀疏数据处理、计算效率优化
    """
    # 深拷贝防止原始数据污染
    df = df.copy(deep=True)
    original_length = len(df)

    try:
        # ===== 输入验证增强 =====
        # 验证必须字段存在性及数据类型
        required_cols = {'close': 'float64', 'volume': 'float64'}
        missing_cols = [col for col, dtype in required_cols.items() if col not in df.columns]
        if missing_cols:
            raise ValueError(f"关键字段缺失: {missing_cols}")

        type_errors = []
        for col, expected_type in required_cols.items():
            if not np.issubdtype(df[col].dtype, np.dtype(expected_type).type):
                type_errors.append(f"{col}应为{expected_type}, 实际为{df[col].dtype}")
        if type_errors:
            raise TypeError(f"字段类型错误:\n" + "\n".join(type_errors))

        # ===== 数据预处理增强 =====
        # 填充策略：限制最大填充步长
        fill_limit = 3  # 允许最多连续填充3个NaN
        df[['close', 'volume']] = df[['close', 'volume']].ffill(limit=fill_limit).bfill(limit=fill_limit)

        # 空值统计与处理
        null_counts = df[['close', 'volume']].isnull().sum()
        if null_counts.any():
            problematic = null_counts[null_counts > 0].to_dict()
            raise ValueError(f"不可填充的空值存在:\n{problematic}")

        # ===== 动态窗口验证 =====
        windows_config = {
            'price': {'series': 'close', 'windows': [5, 10, 20, 30]},
            'volume': {'series': 'volume', 'windows': [5, 10, 20, 30]}
        }
        all_windows = [w for c in windows_config.values() for w in c['windows']]
        max_window = max(all_windows)

        # 提前验证数据长度
        if len(df) < max_window:
            available = len(df)
            required = f"至少需要{max_window}个数据点"
            raise ValueError(f"数据长度不足: {available} < {required}")

        # ===== 计算效率优化 =====
        # 预计算所有滚动窗口
        rolling_objects = {}
        for indicator_type, config in windows_config.items():
            base_series = df[config['series']]
            for w in config['windows']:
                key = f"{indicator_type}_ma{w}"
                rolling_objects[key] = base_series.rolling(
                    window=w, min_periods=w, closed='right'
                )

        # 批量计算均线
        ma_df = pd.DataFrame({
            key: roller.mean() for key, roller in rolling_objects.items()
        })

        # 并行计算偏离度
        def calc_deviation(col):
            ma_series = ma_df[col]
            base = df[windows_config[col.split('_')[0]]['series']]
            return ((base - ma_series) / (ma_series.abs() + 1e-7)) * 100

        deviation_cols = [f"{col}_pct" for col in ma_df.columns]
        deviation_df = pd.concat(
            [calc_deviation(col) for col in ma_df.columns],
            axis=1,
            keys=deviation_cols
        )

        # ===== 数据拼接与修剪 =====
        full_df = pd.concat([df, ma_df, deviation_df], axis=1)

        # 智能修剪无效区间（保留最大窗口后的数据）
        valid_start = max_window - 1
        clean_df = full_df.iloc[valid_start:].copy()

        # ===== 异常值鲁棒处理 =====
        pct_cols = [col for col in deviation_cols if '_pct' in col]
        for col in pct_cols:
            # 基于MAD的异常值处理
            med = clean_df[col].median()
            mad = (clean_df[col] - med).abs().median()
            if mad > 0:
                clean_df[col] = np.clip(
                    clean_df[col],
                    med - 3 * mad,
                    med + 3 * mad
                )

        # 最终完整性检查
        if clean_df.empty:
            loss_ratio = 1.0 - len(clean_df) / original_length
            raise RuntimeError(
                f"数据有效率为0，最大窗口{max_window}修剪后无剩余数据，"
                f"原始数据长度: {original_length}"
            )

        return clean_df

    except Exception as e:
        error_msg = (
            f"技术指标计算失败: {str(e)}\n"
            f"输入数据摘要:\n"
            f"- 长度: {original_length}\n"
            f"- 起止日期: {df.index.min()} ~ {df.index.max()}\n"
            f"- 字段: {df.columns.tolist()}"
        )
        raise RuntimeError(error_msg) from e
